#!/usr/bin/env python3

import sys
import os
import hashlib
import psycopg2
import pdb
import logging
import click
import logging

pg_username='postgres'
pg_password='Post123'
pg_database='postgres'
pg_host='172.32.148.156'
pg_port='1921'
pg_schema='public'

log_file = os.environ.get('LOGDIR', 'logs') + '/log-print-def-pg.log'

print('Purge log file: {}'.format(log_file))
logging.basicConfig(level=logging.INFO,
                    filename=log_file,
                    filemode='w')


@click.group(context_settings=dict(help_option_names=['-h', '--help']))
def main():
    pass

# One chunk is a list of tables
# returns a list of chunks
def get_chunks(table_list, input_file, parallel):
    tablist = []
    if table_list:
        tablist = table_list.split(',')
    elif input_file:
        with open(input_file, 'r', encoding="utf-8") as f:
            while True:
                tabname = f.readline().strip()
                if len(tabname) > 0:
                    tablist.append(tabname)
                else:
                    break;
    else:
        raise Exception('Error: No input tables.')

    chunk_size = int(len(tablist) / parallel)
    if chunk_size <= 0:
        raise Exception('Error: table number: {}, parallel: {}.'.format(len(tablist), parallel))

    chunked_list = [tablist[i:i+chunk_size] for i in range(0, len(tablist), chunk_size)]
    return chunked_list


@main.command('tab', help='Print table names')
def get_tables():
    pg_conn = psycopg2.connect(database=pg_database, user=pg_username, password=pg_password, host=pg_host, port=pg_port)
    pg_cur=pg_conn.cursor()
    pg_cur.arraysize = 100
    pg_cur.execute("""select tablename from pg_tables where schemaname=%s and tableowner=%s""",
                                                   (pg_schema.lower(),pg_username.lower(),))
    res = pg_cur.fetchall()
    for pg_row in res:
        if pg_row[0] != None:
            print(pg_row[0].upper())
    pg_conn.close()


@main.command('seq', help='Print sequence names')
def get_sequences():
    pg_conn = psycopg2.connect(database=pg_database, user=pg_username, password=pg_password, host=pg_host, port=pg_port)
    pg_cur=pg_conn.cursor()
    pg_cur.arraysize = 100
    pg_cur.execute("""select sequencename from pg_sequences where schemaname = %s and sequenceowner = %s""",
                                                    (pg_schema.lower(), pg_username.lower(),))
    res = pg_cur.fetchall()
    for pg_row in res:
        if pg_row[0] != None:
            print(pg_row[0].upper())
    pg_conn.close()


@main.command('tabdef', help='Print defination of each table.', no_args_is_help=True)
@click.option('--input-file', '-i', help='Intput file')
@click.option('--tab-list', '-l', help='Input table list e.g. -l "tb1,tb2,tb3"')
def get_tabdef(input_file, tab_list):
    chunks = get_chunks(tab_list, input_file, 1)
    pg_conn = psycopg2.connect(database=pg_database, user=pg_username, password=pg_password, host=pg_host, port=pg_port)
    pg_cur = pg_conn.cursor()
    for tablename in chunks[0]:
        print('{}:'.format(tablename.upper()))
        pg_cur.execute("""select upper(column_name), data_type, is_nullable from information_schema.columns
                          where table_schema = %s and table_name = %s order by ordinal_position""",
                          (pg_schema.lower(), tablename.lower(),))
        pg_res = pg_cur.fetchall()
        for i in range(len(pg_res)):
            print('    {},{},{}'.format(pg_res[i][0].upper(), pg_res[i][1].upper(), pg_res[i][2].upper()))
        if pg_res == None:
            print('Table {} not exists.'.format(tablename.upper()))
    pg_conn.close()


@main.command('pk', help='Print primary key defination of each table.', no_args_is_help=True)
@click.option('--input-file', '-i', help='Intput file')
@click.option('--tab-list', '-l', help='Input table list e.g. -l "t1,t2,t3"')
def pk_def(input_file, tab_list):
    chunks = get_chunks(tab_list, input_file, 1)
    pg_conn = psycopg2.connect(database=pg_database, user=pg_username, password=pg_password, host=pg_host, port=pg_port)
    pg_cur = pg_conn.cursor()
    for tablename in chunks[0]:
        # pg primary key columns in define order
        pg_cur.execute("""select upper(i.relname) as index_name,
                                 upper(array_to_string(array_agg(a.attname), ',')) as column_names
                          from pg_class t, pg_class i, pg_index ix, pg_attribute a, pg_namespace n
                          where t.oid = ix.indrelid and i.oid = ix.indexrelid
                                and a.attrelid = t.oid and a.attnum = ANY(ix.indkey)
                                and t.relkind = 'r' and t.relnamespace = n.oid
                                and ix.indisprimary = 't'
                                and n.nspname = %s and t.relname = %s
                          group by t.relname, i.relname, ix.indkey
                          order by t.relname, i.relname""",
                          (pg_schema.lower(), tablename.lower(),))
        pg_res = pg_cur.fetchall()
        if pg_res == None:
            print('Table {} not exists.'.format(tablename.upper()))
        else:
            print('{}:'.format(tablename.upper()))
            for pg_row in pg_res:
                print('    {} | {}'.format(pg_row[0].upper(), pg_row[1].upper()))
    pg_conn.close()


@main.command('index', help='Print primary key, index and unique defination of each table.', no_args_is_help=True)
@click.option('--input-file', '-i', help='Intput file')
@click.option('--tab-list', '-l', help='Input table list e.g. -l "t1,t2,t3"')
def index_def(input_file, tab_list):
    chunks = get_chunks(tab_list, input_file, 1)
    pg_conn = psycopg2.connect(database=pg_database, user=pg_username, password=pg_password, host=pg_host, port=pg_port)
    pg_cur = pg_conn.cursor()
    for tablename in chunks[0]:
        # pg index cols in define order
        pg_cur.execute("""select upper(i.relname) as index_name,
                                 upper(array_to_string(array_agg(a.attname), ',')) as column_names
                          from pg_class t, pg_class i, pg_index ix, pg_attribute a, pg_namespace n
                          where t.oid = ix.indrelid and i.oid = ix.indexrelid
                                and a.attrelid = t.oid and a.attnum = ANY(ix.indkey)
                                and t.relkind = 'r' and t.relnamespace = n.oid
                                and n.nspname = %s and t.relname = %s
                          group by t.relname, i.relname, ix.indkey
                          order by t.relname, i.relname""",
                          (pg_schema.lower(), tablename.lower(),))
        pg_res = pg_cur.fetchall()
        if pg_res == None:
            print('Table {} not exists.'.format(tablename.upper()))
        else:
            print('{}:'.format(tablename.upper()))
            for pg_row in pg_res:
                print('    {} | {}'.format(pg_row[0].upper(), pg_row[1].upper()))
    pg_conn.close()



@main.command('fk', help='Print foreign key defination of each table.', no_args_is_help=True)
@click.option('--input-file', '-i', help='Intput file')
@click.option('--tab-list', '-l', help='Input table list e.g. -l "t1,t2,t3"')
def cmp_fk(input_file, tab_list):
    chunks = get_chunks(tab_list, input_file, 1)
    pg_conn = psycopg2.connect(database=pg_database, user=pg_username, password=pg_password, host=pg_host, port=pg_port)
    pg_cur = pg_conn.cursor()
    for tablename in chunks[0]:
        pg_cur.execute("""select conname, 
                                 (select array_to_string(array_agg(a.attname),',') 
                                     from pg_attribute a 
                                     where c.conrelid = a.attrelid and 
                                           a.attnum = any(c.conkey)) fk_cols, 
                                 confrelid::regclass rname, 
                                 (select array_to_string(array_agg(b.attname),',') 
                                     from pg_attribute b 
                                     where c.confrelid = b.attrelid and 
                                           b.attnum = any(c.confkey) ) rk_cols 
                          from pg_constraint c 
                          where conkey is not null and confkey is not null and contype = 'f' and 
                                connamespace = %s::regnamespace::oid and conrelid = %s::regclass::oid
                          order by fk_cols, rname, rk_cols""",
                          (pg_schema.lower(), tablename.lower(),)) 
        pg_res = pg_cur.fetchall()
        if pg_res == None:
            print('Table {} not exists.'.format(tablename.upper()))
        else:
            print('{}:'.format(tablename.upper()))
            for pg_row in pg_res:
                print('    {}({}) references {}({})'.format(pg_row[0].upper(), pg_row[1].upper(), pg_row[2].upper(), pg_row[3].upper()))
    pg_conn.close()

if __name__ == '__main__':
    main(None)

